"""
日志记录工具
提供训练过程的日志记录和保存功能
"""
import json
import pandas as pd
from pathlib import Path
from datetime import datetime
import sys


class Logger:
    """
    训练日志记录器
    
    功能:
        - 记录训练过程到文本文件
        - 保存指标到CSV文件
        - 打印到控制台
    
    参数:
        log_dir: 日志保存目录
        experiment_name: 实验名称
    """
    
    def __init__(self, log_dir, experiment_name):
        self.log_dir = Path(log_dir)
        self.experiment_name = experiment_name
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        # 日志文件
        self.log_file = self.log_dir / f'{experiment_name}_log.txt'
        self.metrics_file = self.log_dir / f'{experiment_name}_metrics.csv'
        
        # 初始化日志
        self._init_log()
    
    def _init_log(self):
        """初始化日志文件"""
        with open(self.log_file, 'w', encoding='utf-8') as f:
            f.write("="*70 + "\n")
            f.write(f"Training Log - {self.experiment_name}\n".center(70))
            f.write("="*70 + "\n")
            f.write(f"Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("="*70 + "\n\n")
        
        print(f"Logger initialized: {self.log_file}")
    
    def log(self, message, print_to_console=True):
        """
        记录一般消息
        
        参数:
            message: 消息内容
            print_to_console: 是否打印到控制台
        """
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        log_message = f"[{timestamp}] {message}"
        
        # 写入文件
        with open(self.log_file, 'a', encoding='utf-8') as f:
            f.write(log_message + "\n")
        
        # 打印到控制台
        if print_to_console:
            print(log_message)
    
    def log_epoch(self, epoch, train_loss, train_metrics, val_loss, val_metrics):
        """
        记录每个epoch的信息
        
        参数:
            epoch: epoch编号
            train_loss: 训练损失
            train_metrics: 训练指标字典
            val_loss: 验证损失
            val_metrics: 验证指标字典
        """
        # 写入文本日志
        with open(self.log_file, 'a', encoding='utf-8') as f:
            f.write(f"\nEpoch {epoch}\n")
            f.write("-" * 70 + "\n")
            f.write(f"  Train - Loss: {train_loss:.4f}, MAE: {train_metrics['MAE']:.4f}, "
                   f"RMSE: {train_metrics['RMSE']:.4f}, R²: {train_metrics['R²']:.4f}\n")
            f.write(f"  Val   - Loss: {val_loss:.4f}, MAE: {val_metrics['MAE']:.4f}, "
                   f"RMSE: {val_metrics['RMSE']:.4f}, R²: {val_metrics['R²']:.4f}\n")
    
    def log_config(self, config):
        """
        记录配置信息
        
        参数:
            config: 配置对象
        """
        with open(self.log_file, 'a', encoding='utf-8') as f:
            f.write("\nConfiguration:\n")
            f.write("-" * 70 + "\n")
            config_dict = config.to_dict() if hasattr(config, 'to_dict') else vars(config)
            for key, value in config_dict.items():
                f.write(f"  {key}: {value}\n")
            f.write("\n")
    
    def save_history(self, history):
        """
        保存训练历史到CSV
        
        参数:
            history: 训练历史字典，包含:
                - train_loss: 训练损失列表
                - val_loss: 验证损失列表
                - train_metrics: 训练指标列表
                - val_metrics: 验证指标列表
                - learning_rates: 学习率列表
        """
        records = []
        
        for i in range(len(history['train_loss'])):
            record = {
                'epoch': i + 1,
                'train_loss': history['train_loss'][i],
                'val_loss': history['val_loss'][i],
                'learning_rate': history['learning_rates'][i] if i < len(history['learning_rates']) else 0,
            }
            
            # 添加训练指标
            if i < len(history['train_metrics']):
                train_m = history['train_metrics'][i]
                for k, v in train_m.items():
                    record[f'train_{k}'] = v
            
            # 添加验证指标
            if i < len(history['val_metrics']):
                val_m = history['val_metrics'][i]
                for k, v in val_m.items():
                    record[f'val_{k}'] = v
            
            records.append(record)
        
        # 保存到CSV
        df = pd.DataFrame(records)
        df.to_csv(self.metrics_file, index=False, encoding='utf-8-sig')
        
        print(f"Training history saved to {self.metrics_file}")
        print(f"  Total epochs: {len(df)}")
        print(f"  Metrics columns: {len(df.columns)}")
    
    def log_final_results(self, best_epoch, best_metric, total_time=None):
        """
        记录最终训练结果
        
        参数:
            best_epoch: 最佳epoch
            best_metric: 最佳指标值
            total_time: 总训练时间（秒）
        """
        with open(self.log_file, 'a', encoding='utf-8') as f:
            f.write("\n" + "="*70 + "\n")
            f.write("Training Completed\n".center(70))
            f.write("="*70 + "\n")
            f.write(f"Best Epoch: {best_epoch}\n")
            f.write(f"Best R²: {best_metric:.4f}\n")
            if total_time:
                f.write(f"Total Time: {total_time:.2f} seconds ({total_time/3600:.2f} hours)\n")
            f.write(f"Completed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("="*70 + "\n")


class TensorBoardLogger:
    """
    TensorBoard日志记录器（可选）
    
    需要安装: pip install tensorboard
    """
    
    def __init__(self, log_dir, experiment_name):
        try:
            from torch.utils.tensorboard import SummaryWriter
            self.writer = SummaryWriter(log_dir=f'{log_dir}/{experiment_name}')
            self.enabled = True
            print(f"TensorBoard logging enabled: {log_dir}/{experiment_name}")
        except ImportError:
            print("Warning: tensorboard not installed. TensorBoard logging disabled.")
            self.enabled = False
    
    def log_scalar(self, tag, value, step):
        """记录标量值"""
        if self.enabled:
            self.writer.add_scalar(tag, value, step)
    
    def log_scalars(self, main_tag, tag_scalar_dict, step):
        """记录多个标量值"""
        if self.enabled:
            self.writer.add_scalars(main_tag, tag_scalar_dict, step)
    
    def log_metrics(self, metrics, step, prefix=''):
        """记录指标字典"""
        if self.enabled:
            for key, value in metrics.items():
                self.writer.add_scalar(f'{prefix}/{key}', value, step)
    
    def close(self):
        """关闭writer"""
        if self.enabled:
            self.writer.close()


if __name__ == "__main__":
    # 测试Logger
    print("Testing Logger...")
    
    import tempfile
    import shutil
    
    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    
    try:
        # 创建logger
        logger = Logger(temp_dir, 'test_experiment')
        
        # 记录消息
        logger.log("Test message 1")
        logger.log("Test message 2", print_to_console=False)
        
        # 记录epoch
        train_metrics = {'MAE': 1.5, 'RMSE': 2.0, 'R²': 0.85}
        val_metrics = {'MAE': 1.8, 'RMSE': 2.2, 'R²': 0.80}
        logger.log_epoch(1, 0.5, train_metrics, 0.6, val_metrics)
        
        # 保存历史
        history = {
            'train_loss': [0.5, 0.4, 0.3],
            'val_loss': [0.6, 0.5, 0.4],
            'train_metrics': [train_metrics] * 3,
            'val_metrics': [val_metrics] * 3,
            'learning_rates': [0.001, 0.0008, 0.0006]
        }
        logger.save_history(history)
        
        # 记录最终结果
        logger.log_final_results(3, 0.85, 3600)
        
        print("\n✓ Logger test passed!")
        
    finally:
        # 清理临时目录
        shutil.rmtree(temp_dir)